# -*- coding: utf-8 -*-
'''
Created on 15.12.2019

@author: yu03
'''
import numpy as np
import os
import re
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import glob
from mpl_toolkits.mplot3d import Axes3D
from PyUeye_Unified.Cross_200line_SAAC import folder_path, np_result_names, hor_index, ver_index, hor_lines, ver_lines
import sys
from FFT_Interpolation import FFT_interpolation_nonlinearity_compare
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter

def running_mean(x, N):
    cumsum = np.cumsum(np.insert(x, 0, 0)) 
    return (cumsum[N:]-cumsum[:-N])/float(N)

i = 0
fs_cam = 25
hor_angle_avr_set = []
ver_angle_avr_set = []
hor_length_avr_set = []
ver_length_avr_set = []
time_set = []
''' 
    for all groups
''' 
for np_result in np_result_names:
    ''' 
        reading .npy file
    ''' 
    file_name = np_result.split('\\')[-1] ### x_lines.py
#     file_name = re.findall(r"openmode\\(.+).npy", np_result)[0]
    print(file_name)
    i += 1
    f = open(np_result, 'rb')
    f_size = os.fstat(f.fileno()).st_size
    ''' 
        put 4-D results in "lines"
    ''' 
    lines = []
    time_sequence = []
    while f.tell() < f_size:
        line = np.load(f, allow_pickle=True)
        lines.append(line)
    print('%i: %s, %.1fMB'%(i, file_name, f_size/1e6))
    time_sequence = lines[4]
    lines = np.array(lines[0:4])
    print(np.shape(lines), 'Channle, lines, frames:')
    frame_num = np.shape(lines)[2]
    line_num = np.shape(lines)[1]
    
#     print(np.shape(lines))
    hor_angle_set, ver_angle_set, hor_length_set, ver_length_set = lines
    time_set.append(time_sequence)
    
    ''' 
        making average
    ''' 
    averaged_results = []
    
#     index_line = 0
#     hor_angle_avr = hor_angle_set[index_line]
#     ver_angle_avr = ver_angle_set[index_line]
#     hor_length_avr = hor_length_set[index_line]
#     ver_length_avr = ver_length_set[index_line]
    
    hor_angle_avr = np.mean(hor_angle_set, axis=0)
    ver_angle_avr = np.mean(ver_angle_set, axis=0)
    hor_length_avr = np.mean(hor_length_set, axis=0)
    ver_length_avr = np.mean(ver_length_set, axis=0)
    
    hor_angle_avr_set.append(hor_angle_avr)
    ver_angle_avr_set.append(ver_angle_avr)
    hor_length_avr_set.append(hor_length_avr)
    ver_length_avr_set.append(ver_length_avr)

time_set = np.array(time_set)
time_set = np.hstack(time_set)  
time_set = time_set + 0.3

hor_angle_avr_set = np.array(hor_angle_avr_set)
ver_angle_avr_set = np.array(ver_angle_avr_set)
hor_length_avr_set = np.array(hor_length_avr_set)
ver_length_avr_set = np.array(ver_length_avr_set)

hor_angle_avr_set = np.hstack(hor_angle_avr_set)  
ver_angle_avr_set = np.hstack(ver_angle_avr_set)
hor_length_avr_set = np.hstack(hor_length_avr_set)
ver_length_avr_set = np.hstack(ver_length_avr_set)

hor_angle_avr_set = (hor_angle_avr_set - hor_angle_avr_set[0]) * 0.97 + hor_angle_avr_set[0]
ver_angle_avr_set = (ver_angle_avr_set - ver_angle_avr_set[0]) * 0.97 + ver_angle_avr_set[0]


'''
    write
'''
with open(r'E:\3DoF_Interferometer_SAAC\Drift\results_2line_060_-150001\lines_averaged\lines_averaged_result.npy', 'wb') as file_name:
    np.savez(file_name, 
             hor_angle_avr_set=hor_angle_avr_set,ver_angle_avr_set=ver_angle_avr_set,
             hor_length_avr_set=hor_length_avr_set,ver_length_avr_set=ver_length_avr_set,
             time_set=time_set)


hor_angle_avr_set = running_mean(hor_angle_avr_set, 200)
ver_angle_avr_set = running_mean(ver_angle_avr_set, 200)
hor_length_avr_set = running_mean(hor_length_avr_set, 200)
ver_length_avr_set = running_mean(ver_length_avr_set, 200)
time_set = time_set[100:-99]

print(len(hor_angle_avr_set), hor_angle_avr_set[0], hor_angle_avr_set[-1])
print(len(time_set), time_set[0], time_set[-1])

fig = plt.figure('Raw Data')
plt.gcf().set_size_inches(18,9)

ax1 = fig.add_subplot(2, 2, 1)
ax2 = fig.add_subplot(2, 2, 2)
ax3 = fig.add_subplot(2, 2, 3)
ax4 = fig.add_subplot(2, 2, 4)

ax1.plot(time_set, hor_angle_avr_set, color='black', marker=' ', markersize=1, label='horizontal angle')
ax2.plot(time_set, hor_length_avr_set, color='black', marker=' ', markersize=1, label='horizontal length')
ax3.plot(time_set, ver_angle_avr_set, color='black', label='vertical angle')
ax4.plot(time_set, ver_length_avr_set, color='black', label='vertical length')

ax1.title.set_text('Horizontal Tilt')
ax2.title.set_text('Horizontal Length')
ax3.title.set_text('Vertical Tilt')
ax4.title.set_text('Vertical Length')
   
ax1.set_ylabel('Horizontal Tilt \ urad')
ax1.set_xlabel('Timestamp')
   
ax2.set_ylabel('Horizontal Length \ nm')
ax2.set_xlabel('Frame Num.')
   
ax3.set_ylabel('Vertical Tilt \ urad')
ax3.set_xlabel('Timestamp')
   
ax4.set_ylabel('Vertical Length \ nm')
ax4.set_xlabel('Frame Num.')

ax1.grid(which='both', axis='both')
ax2.grid(which='both', axis='both')
ax3.grid(which='both', axis='both')
ax4.grid(which='both', axis='both')

ax1.legend()
ax2.legend()
ax3.legend()
ax4.legend()

plt.tight_layout()
plt.show()



